import pandas as pd
import numpy as np

def selectFrame(input, gap=3, selecteFrameNum=1):  # 抽帧操作 减少计算成本
    if selecteFrameNum >= gap:
        print("参数设置错误，selecteFrameNum必须要比gap小。")
        return
    output = []
    count = 0
    selecteFrameCount = 0
    for i in range(0, len(input)):
        if count == gap:
            count = 0
            selecteFrameCount = 0
        if selecteFrameNum > selecteFrameCount:
            output.append(input[i])
        count += 1
        selecteFrameCount += 1
    return output


'''
select_person需要传入一个参数：
    @params:person_code 待选择人物的编号，以0开始
'''


def select_person(input, params):
    person_code = params["person_code"]
    print("选择人物操作完成，人物id为：{}".format(person_code))
    return list(map(lambda x: [x[person_code]], input))


'''
    @:key_points_indexes: 待选择的关键点序号列表，具体参考mmpose的介绍
    @:show_bbox: 是否需要bbox，传入一个布尔值
'''


def extract_designative_key_points(input, params):
    key_points_indexes = params["key_points_indexes"]
    show_bbox = params["show_bbox"]

    def key_points_filter(enumerate_item):
        if key_points_indexes == "all":
            return True
        index = enumerate_item[0]
        if index in key_points_indexes:
            return True
        return False

    for frame_data in input:  # 某时间
        for obj_data in frame_data:  # 某个人
            if not show_bbox:
                obj_data.bbox = None
            if key_points_indexes == "all":
                pass
            else:
                obj_data.key_points = list(
                    map(lambda x: x[1], filter(key_points_filter, enumerate(obj_data.key_points))))
    print("抽取指定特征点工作完成")
    return input


def smooth_key_points(input, params):
    # 用滑动窗口的方式，使关键点的变化变得平滑，减少异常抖动和异常帧的影响
    EM_num = params["EM_num"]
    cache_map = {}
    # 构建cache_map
    for frame_data in input:  # 某时间
        for obj_data_index, obj_data in enumerate(frame_data):  # 某个人
            for key_point_index, key_point in enumerate(obj_data.key_points):
                x_point = key_point[0]
                y_point = key_point[1]
                if obj_data_index not in cache_map:
                    cache_map[obj_data_index] = {}
                if key_point_index not in cache_map[obj_data_index]:
                    cache_map[obj_data_index][key_point_index] = []
                cache_map[obj_data_index][key_point_index].append((x_point, y_point))
    # 执行平滑操作
    for frame_index, frame_data in enumerate(input):
        for obj_index, obj_data in enumerate(frame_data):
            for key_point_index, key_point in enumerate(obj_data.key_points):
                if frame_index + 1 <= EM_num:
                    continue
                else:
                    new_x, new_y = _calc_smooth_val(cache_map, obj_index, key_point_index, frame_index, EM_num)
                    input[frame_index][obj_index].key_points[key_point_index] = (
                        new_x, new_y, input[frame_index][obj_index].key_points[key_point_index][2]
                    )
    print("平滑关键点操作完成")
    return input


def _calc_smooth_val(cache_map, obj_index, key_point_index, frame_index, EM_num):
    sub_list = cache_map[obj_index][key_point_index][(frame_index - EM_num + 1): (frame_index + 1)]
    sub_list_df = pd.DataFrame(sub_list)
    sub_list_df_mean = sub_list_df.mean(axis=0)
    # print("sub_list_df_mean[0]", sub_list_df_mean[0], " sub_list_df_mean[1]", sub_list_df_mean[1])
    return sub_list_df_mean[0], sub_list_df_mean[1]

